from langchain.vectorstores import Chroma
# from langchain_chroma import Chroma

from config.embedding_config import get_openai_embeddings_xin, get_openai_embeddings_local

# 1. 配置环境
PERSIST_DIR = "./chroma_db"  # 持久化目录

import sentence_transformers
print(sentence_transformers.__version__)

# 3. 初始化嵌入模型
embeddings = get_openai_embeddings_local()

# 加载已有向量库
vector_db = Chroma(
    persist_directory=PERSIST_DIR,
    embedding_function=embeddings
)

query = "如何联系客户服务"

# 相似性搜索（返回原始文档）
results = vector_db.similarity_search(
    query=query,  # 支持直接输入文本
    k=3,  # 返回结果数量
    include=["documents", "metadatas", "distances"] # 包含元数据
    # filter={"category": "technology"}  # 元数据过滤
)

print(results)

for i, doc in enumerate(results):
    print(f"结果 {i+1}:")
    print(doc.page_content)
    print("元数据:", doc.metadata)
    print("="*50)

# print(f"✅ 已保存 {len(texts)} 个文本块到 {PERSIST_DIR}")
